/*
这个程序主要进行图染色，要染色的图包括5个节点，邻接矩阵如下
| 1 1 0 0 0 |
| 1 1 1 1 0 |
| 0 1 1 0 1 |
| 0 1 0 1 0 |
| 0 0 1 0 1 |
这个图最少两个颜色，0、2、3一个颜色，1、4一个颜色
https://stackoverflow.com/questions/18027278/improving-the-solution-of-sparse-linear-systems/33483695
*/

#include<stdio.h>
#include<stdlib.h>
#include<string.h>
#include<cusparse.h>

int main(){
    int size = 10000;
    //首先我们使用coo格式来创建邻接矩阵，然后转化为csr，一共13个点
    int nz = 10000;
    //创建三个CPU端数组，分别保存行、列和值。
    int* row_coo_host = 0;
    int* col_host = 0;
    float* val_host = 0;

    //在CPU端申请空间
    cudaMallocHost((void **)&row_coo_host, nz * sizeof(int));
    cudaMallocHost((void **)&col_host, nz * sizeof(int));
    cudaMallocHost((void **)&val_host, nz * sizeof(float));

    //值全是1
    //遍历变量
    int i;
    for(i = 0; i < nz; i++){
        val_host[i] = 1;
    }

    //矩阵构成
    // row_coo_host[0] = 0;
    // col_host[0] = 0;
    // row_coo_host[1] = 0;
    // col_host[1] = 1;
    // row_coo_host[2] = 1;
    // col_host[2] = 0;
    // row_coo_host[3] = 1;
    // col_host[3] = 1;
    // row_coo_host[4] = 1;
    // col_host[4] = 2;
    // row_coo_host[5] = 1;
    // col_host[5] = 3;
    // row_coo_host[6] = 2;
    // col_host[6] = 1;
    // row_coo_host[7] = 2;
    // col_host[7] = 2;
    // row_coo_host[8] = 2;
    // col_host[8] = 4;
    // row_coo_host[9] = 3;
    // col_host[9] = 1;
    // row_coo_host[10] = 3;
    // col_host[10] = 3;
    // row_coo_host[11] = 4;
    // col_host[11] = 2;
    // row_coo_host[12] = 4;
    // col_host[12] = 4;

    //我们构造一个100*100的单位阵
    for(i = 0; i < nz; i++){
        row_coo_host[i] = i;
        col_host[i] = i;
    }


    //实际上coo到csr只有存储行号的数组的变化的，所以cusparseXcoo2csr函数的返回值
    //是新的行号序列，我们在GPU上申请一段新的空间，来存储压缩版的行号，压缩版的行号需要
    //nz + 1个空间来存储每一行的的起始元素在col_host和val_host中的索引
    int* row_host = 0;
    cudaMallocHost((void **)&row_host, (size + 1) * sizeof(int));

    //在设备端申请空间，对于这里我们只需要申请两个用来存储行索引的空间
    int* row_coo_device = 0;
    int* row_device = 0;
    cudaMalloc((void **)&row_coo_device, nz * sizeof(int));
    cudaMalloc((void **)&row_device, (size + 1) * sizeof(int));

    //将row_coo_host的数据传入
    cudaMemcpy(row_coo_device, row_coo_host, nz * sizeof(int), cudaMemcpyHostToDevice);
    
    //创造句柄
    cusparseHandle_t handle;
    cusparseCreate(&handle);

    //调用cusparseXcoo2csr
    cusparseXcoo2csr(handle, row_coo_device, nz, size, row_device, CUSPARSE_INDEX_BASE_ZERO);

    //这里将数据传回CPU
    cudaMemcpy(row_host, row_device, (size+1) * sizeof(int), cudaMemcpyDeviceToHost);

    //打印数组
    //迭代
    for(i = 0 ; i < size + 1; i++){
        printf("%d,", row_host[i]);
    }

    printf("\n");

    //析构两个设备端数组
    cudaFree(row_coo_device);
    cudaFree(row_device);

    //row_host、col_host、val_host分别分别存储CSR存储的行、列、值，下面我们进行染色
    /*
    使用的接口：
    cusparseScsrcolor(cusparseHandle_t handle, int m, int nnz,
             const cusparseMatDescr_t descrA, const float *csrValA,
             const int *csrRowPtrA, const int *csrColIndA,
             const float *fractionToColor, int *ncolors, int *coloring,
             int *reordering,cusparseColorInfo_t info);
    */
    //我们需要构建的是后面的4个参量，首先是info，这个参数实际上是用来在各个CUDA API之间传递信息的
    //虽然我们只调用了一次，但是还是要创建一下这个变量
    cusparseColorInfo_t info;
    cusparseCreateColorInfo(&info);

    //用来存储颜色的数组，我们现在不知道这里面会得到什么，最坏情况，颜色和点的数量是一致的，我们需要按照最坏情况来申请数组
    //我们在gpu上申请数组
    int *coloring_device;
    cudaMalloc((void **)&coloring_device , size * sizeof(int));
    cudaMemset(coloring_device, 0, size * sizeof(int));

    //申请reordering，我们现在还不知道这个能够得到什么，我们按照点的数量来给他申请空间
    int *reordering_device;
    cudaMalloc((void **)&reordering_device , size * sizeof(int));
    cudaMemset(reordering_device, 0, size * sizeof(int));

    //这个是需要严谨染色的节点染色的比例，我们设定为1，让所有的节点全部接受严谨染色
    float fractionToColor = 1;
    //颜色的数量
    int ncolors = 0;
    //上面这两个变量都是单个变量，不用在显存上申请空间，并且来回传输就可以通过指针直接修改

    //初始化矩阵的描述信息，cusparseMatDescr_t这个结构体通常保存了矩阵的索引标准和矩阵类型等信息。
    cusparseMatDescr_t descr;
    cusparseCreateMatDescr(&descr);

    //设定为普通矩阵，从0开始索引
    cusparseSetMatType(descr, CUSPARSE_MATRIX_TYPE_GENERAL);
    cusparseSetMatIndexBase(descr, CUSPARSE_INDEX_BASE_ZERO);

    //将矩阵A全部传到GPU端
    int* col_device = 0;
    float* val_device = 0;
    row_device = 0;

    //申请空间
    cudaMalloc((void **)&col_device, nz * sizeof(int));
    cudaMalloc((void **)&val_device, nz * sizeof(float));
    cudaMalloc((void **)&row_device, (size + 1) * sizeof(int));

    //数据拷贝
    cudaMemcpy(col_device, col_host, nz * sizeof(int), cudaMemcpyHostToDevice);
    cudaMemcpy(val_device, val_host, nz * sizeof(float), cudaMemcpyHostToDevice);
    cudaMemcpy(row_device, row_host, (size + 1) * sizeof(int), cudaMemcpyHostToDevice);
    
    cusparseScsrcolor(handle, size, nz, descr, val_device, row_device, col_device, &fractionToColor, 
        &ncolors, coloring_device, reordering_device, info);
    
    //将数据拷贝回来，主要是要取回coloring_device和reordering_device
    //在CPU端申请空间
    int* coloring_host = 0;
    int* reordering_host = 0;
    cudaMallocHost((void **)&coloring_host, size * sizeof(int));
    cudaMallocHost((void **)&reordering_host, size * sizeof(int));

    //将数据拷贝回CPU端
    cudaMemcpy(coloring_host, coloring_device, size * sizeof(int), cudaMemcpyDeviceToHost);
    cudaMemcpy(reordering_host, reordering_device, size * sizeof(int), cudaMemcpyDeviceToHost);

    printf("coloring_host:");
    //打印CPU端数据
    for(i = 0 ; i < size ; i++){
        printf("%d " , coloring_host[i]);
    }
    printf("\n");

    printf("reordering_host:");
    for(i = 0 ; i < size ; i++){
        printf("%d " , reordering_host[i]);
    }
    printf("\n");
    
    //析构GPU变量
    cudaFree(col_device);
    cudaFree(val_device);
    cudaFree(row_device);
    
    //析构所有的CPU变量
    cudaFreeHost(coloring_host);
    cudaFreeHost(reordering_host);
    cudaFreeHost(row_coo_host);
    cudaFreeHost(col_host);
    cudaFreeHost(val_host);
    cudaFreeHost(row_host);
}